from fastapi import FastAPI, HTTPException, Form
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import pymysql
from typing import List, Any, Dict, Optional
import json
import logging
import base64
import urllib.parse
import os
os.add_dll_directory(r'C:\software\ntx64_odbc_cli\clidriver\bin')
try:
    # DB2 驱动 (需要安装 ibm_db 和 ibm_db_dbi)
    import ibm_db_dbi
except Exception:
    ibm_db_dbi = None

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="数据库查询API", description="用于查询MySQL数据库的FastAPI应用", version="1.0.0")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 数据库配置
DB_CONFIG = {
    "host": "192.168.8.222",
    "port": 3306,
    "user": "root",
    "password": "123456",
    "database": "fms",
    "charset": "utf8mb4"
}

# DB2 数据库配置（审计相关数据源）
DB2_CONFIG = {
    "host": "192.168.8.100",
    "port": 50000,
    "user": "db2inst1",
    "password": "dgls",
    "database": "amsesdb",  
}

class QueryRequest(BaseModel):
    sql: str

class QueryResponse(BaseModel):
    isselect: bool
    totalcount: int
    pageIndex: int
    datas: List[List[Any]]

def get_db_connection():
    """获取数据库连接"""
    try:
        connection = pymysql.connect(**DB_CONFIG)
        return connection
    except Exception as e:
        logger.error(f"数据库连接失败: {e}")
        raise HTTPException(status_code=500, detail=f"数据库连接失败: {str(e)}")

def get_db2_connection(db_name: Optional[str] = None):
    """获取 DB2 数据库连接，可指定数据库名覆盖配置"""
    if ibm_db_dbi is None:
        raise HTTPException(status_code=500, detail="DB2 驱动未安装，请安装 ibm_db 与 ibm_db_dbi")
    try:
        database = (db_name or DB2_CONFIG['database'] or '').strip()
        dsn = (
            f"DATABASE={database};"
            f"HOSTNAME={DB2_CONFIG['host']};"
            f"PORT={DB2_CONFIG['port']};"
            f"PROTOCOL=TCPIP;"
            f"UID={DB2_CONFIG['user']};"
            f"PWD={DB2_CONFIG['password']}"
        )
        connection = ibm_db_dbi.connect(dsn, '', '')
        return connection
    except Exception as e:
        logger.error(f"DB2 数据库连接失败: {e}")
        raise HTTPException(status_code=500, detail=f"DB2 数据库连接失败: {str(e)}")

def is_select_query(sql: str) -> bool:
    """判断是否为SELECT查询"""
    sql_stripped = sql.strip().upper()
    return sql_stripped.startswith('SELECT')

@app.post("/coolv/edatasource/jdbcmgr.do", response_model=QueryResponse)
async def execute_query(sql: str = Form(), jdbcName: str = Form(default="coolv"), database: str = Form(default="")):
    """
    执行SQL查询并返回格式化的结果
    """
    connection = None
    try:
        # 解码SQL: 先Base64解码，再URL解码
        try:
            # 1. Base64解码
            base64_decoded = base64.b64decode(sql).decode('utf-8')
            logger.info(f"Base64解码后: {base64_decoded}")
            
            # 2. URL解码
            decoded_sql = urllib.parse.unquote(base64_decoded).strip()
            print(f"URL解码后(最终SQL): {decoded_sql}")
            
        except Exception as decode_error:
            logger.error(f"SQL解码失败: {decode_error}")
            # 如果解码失败，尝试直接使用原始SQL
            decoded_sql = sql.strip()
            logger.info(f"使用原始SQL: {decoded_sql}")
        
        # 获取数据库连接（根据 jdbcName 分流：coolv -> MySQL，amsedb/db2 -> DB2）
        jdbc = (jdbcName or '').strip().lower()
        if jdbc in ['db2', 'amsesdb']:
            connection = get_db2_connection(database)
        else:
            connection = get_db_connection()

        cursor = connection.cursor()

        cursor.execute(decoded_sql)
        
        # 判断是否为SELECT查询
        is_select = is_select_query(decoded_sql)
        
        if is_select:
            # 获取查询结果
            results = cursor.fetchall()
            
            # 转换为列表格式
            datas = []
            for row in results:
                # 将每一行转换为字符串列表
                row_data = []
                for item in row:
                    if item is None:
                        row_data.append("")
                    else:
                        row_data.append(str(item))
                datas.append(row_data)
            
            response = QueryResponse(
                isselect=True,
                totalcount=len(datas),
                pageIndex=0,
                datas=datas
            )
        else:
            # 非SELECT查询（INSERT, UPDATE, DELETE等）
            affected_rows = cursor.rowcount
            connection.commit()
            
            response = QueryResponse(
                isselect=False,
                totalcount=affected_rows,
                pageIndex=0,
                datas=[]
            )
        
        return response
        
    except pymysql.Error as e:
        logger.error(f"数据库错误: {e}")
        raise HTTPException(status_code=400, detail=f"SQL执行错误: {str(e)}")
    except Exception as e:
        logger.error(f"服务器错误: {e}")
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {str(e)}")
    finally:
        if connection:
            connection.close()

@app.post("/audit/db2/jdbcmgr.do", response_model=QueryResponse)
async def execute_db2_query(sql: str = Form(), database: str = Form(default="")):
    """
    执行 DB2 SQL 查询并返回格式化的结果
    """
    connection = None
    try:
        connection = get_db2_connection(database)
        cursor = connection.cursor()

        # 解码SQL: 先Base64解码，再URL解码
        try:
            base64_decoded = base64.b64decode(sql).decode('utf-8')
            logger.info(f"DB2 Base64解码后: {base64_decoded}")
            decoded_sql = urllib.parse.unquote(base64_decoded).strip()
            print(f"DB2 URL解码后(最终SQL): {decoded_sql}")
        except Exception as decode_error:
            logger.error(f"DB2 SQL解码失败: {decode_error}")
            decoded_sql = sql.strip()

        cursor.execute(decoded_sql)

        is_select = is_select_query(decoded_sql)

        if is_select:
            rows = cursor.fetchall()
            datas: List[List[str]] = []
            for row in rows:
                row_data: List[str] = []
                for item in row:
                    if item is None:
                        row_data.append("")
                    else:
                        row_data.append(str(item))
                datas.append(row_data)

            response = QueryResponse(
                isselect=True,
                totalcount=len(datas),
                pageIndex=0,
                datas=datas
            )
        else:
            affected_rows = cursor.rowcount
            connection.commit()
            response = QueryResponse(
                isselect=False,
                totalcount=affected_rows,
                pageIndex=0,
                datas=[]
            )

        return response
    except Exception as e:
        logger.error(f"DB2 执行错误: {e}")
        raise HTTPException(status_code=500, detail=f"DB2 执行错误: {str(e)}")
    finally:
        if connection:
            try:
                connection.close()
            except Exception:
                pass

@app.get("/")
async def root():
    """根路径，返回API信息"""
    return {
        "message": "数据库查询API服务",
        "version": "1.0.0",
        "endpoints": {
            "POST /coolv/edatasource/jdbcmgr.do": "执行SQL查询",
            "GET /docs": "API文档"
        }
    }



if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8081)
